#ifndef GLOBAL_AVG_POOL_N_D_H
#define GLOBAL_AVG_POOL_N_D_H

#include "kernel_operator.h"

namespace GlobalAvgPool {
using namespace AscendC;

constexpr int32_t FP16_BUFFER_NUM = 1;
constexpr int32_t FP32_BUFFER_NUM = 2;
constexpr int32_t FP16_TEMP_BUF_COUNT = 4;
constexpr int32_t FP32_TEMP_BUF_COUNT = 1;

constexpr int32_t BYTE_BLOCK = 32;

constexpr float POS_ONE = 1.0;

template <typename T>
class GlobalAvgPoolND {
public:
    __aicore__ inline GlobalAvgPoolND() = default;
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, const GlobalAvgPoolTilingData* __restrict tilingData);
    __aicore__ inline void Process();

private:
    template <typename T1, typename T2>
    __aicore__ inline T1 CeilDiv(T1 a, T2 b) {
        T1 bTemp(b);
        return bTemp == 0 ? a : (a + bTemp - 1) / bTemp;
    };

    template <typename T1, typename T2>
    __aicore__ inline T1 CeilAlignA2B(T1 a, T2 b) {
        T1 bTemp(b);
        return bTemp == 0 ? a : CeilDiv(a, bTemp) * bTemp;
    };

    __aicore__ inline void CopyIn(int64_t gmOffset, int64_t dataCount);
    __aicore__ inline void Compute(int64_t dataCount);

private:
    TPipe pipe;

#if ORIG_DTYPE_X == DT_FLOAT
    TQue<QuePosition::VECIN, FP32_BUFFER_NUM> xQue;
#else
    TQue<QuePosition::VECIN, FP16_BUFFER_NUM> xQue;
#endif
    TBuf<QuePosition::VECCALC> tempValBuf;
    GlobalTensor<T> xGM, yGM;
    LocalTensor<float> tempValLT;
    int64_t blockIdx = 0;
    uint64_t perBlockCount = 0;

    // tiling params
    uint64_t ubMaxProcCount = 0;
    uint64_t totalDataCount = 0;
    uint64_t ncSize = 0;
    uint64_t hwSize = 0;
    uint64_t loopTime = 0;
    uint64_t tailCount = 0;
};

template <typename T>
__aicore__ inline void GlobalAvgPoolND<T>::Init(GM_ADDR x, GM_ADDR y,
                                                const GlobalAvgPoolTilingData* __restrict tilingData) {
    blockIdx = GetBlockIdx();
    perBlockCount = BYTE_BLOCK / sizeof(T);
    ubMaxProcCount = tilingData->ubMaxProcCount;
    totalDataCount = tilingData->totalDataCount;
    ncSize = tilingData->ncSize;
    hwSize = tilingData->hwSize;
    loopTime = tilingData->loopTime;
    tailCount = tilingData->tailCount;

    xGM.SetGlobalBuffer((__gm__ T*)x, totalDataCount);
    yGM.SetGlobalBuffer((__gm__ T*)y, ncSize);

    uint64_t singleBufferSize = ubMaxProcCount * sizeof(T);
#if ORIG_DTYPE_X == DT_FLOAT
    pipe.InitBuffer(xQue, FP32_BUFFER_NUM, singleBufferSize);
    pipe.InitBuffer(tempValBuf, singleBufferSize * FP32_TEMP_BUF_COUNT);
#else
    pipe.InitBuffer(xQue, FP16_BUFFER_NUM, singleBufferSize);
    pipe.InitBuffer(tempValBuf, singleBufferSize * FP16_TEMP_BUF_COUNT);
#endif

    tempValLT = tempValBuf.Get<float>();
}

template <typename T>
__aicore__ inline void GlobalAvgPoolND<T>::Process() {
    int64_t gmOffset = 0;
    float sumValue = 0;
    for (int64_t i = 0; i < ncSize; i++) {
        sumValue = 0;
        for (int64_t j = 0; j < loopTime; j++) {
            CopyIn(gmOffset, ubMaxProcCount);
            Compute(ubMaxProcCount);
            sumValue += tempValLT.GetValue(0);
            gmOffset += ubMaxProcCount;
        }

        if (tailCount) {
            int64_t tailAlignCount = CeilAlignA2B(tailCount, perBlockCount);
            CopyIn(gmOffset, tailAlignCount);
            Compute(tailCount);
            sumValue += tempValLT.GetValue(0);
            gmOffset += tailCount;
        }
        yGM.SetValue(i, sumValue / hwSize);
    }
}

template <typename T>
__aicore__ inline void GlobalAvgPoolND<T>::CopyIn(int64_t gmOffset, int64_t dataCount) {
    LocalTensor<T> xInLT = xQue.AllocTensor<T>();
    DataCopy(xInLT, xGM[gmOffset], dataCount);
    xQue.EnQue(xInLT);
}

template <typename T>
__aicore__ inline void GlobalAvgPoolND<T>::Compute(int64_t dataCount) {
#if ORIG_DTYPE_X == DT_FLOAT
    LocalTensor<float> xLT = xQue.DeQue<float>();
#else
    LocalTensor<half> xLTFloat16 = xQue.DeQue<half>();
    LocalTensor<float> xLT = tempValLT[ubMaxProcCount];
    Cast(xLT, xLTFloat16, RoundMode::CAST_NONE, dataCount);
#endif
    LocalTensor<float> t1 = tempValLT;
    ReduceSum(t1, xLT, t1, dataCount);
#if ORIG_DTYPE_X == DT_FLOAT
    xQue.FreeTensor(xLT);
#else
    xQue.FreeTensor(xLTFloat16);
#endif
}

}  // namespace GlobalAvgPool

#endif  // GLOBAL_AVG_POOL_N_D_H